#!/usr/bin/env python
"""

    locus.py
    [--log_file PATH]
    [--verbose]

"""

################################################################################
#
#   locus
#
#
#   Copyright (c) 12/15/2009 Leo Goodstadt
#   
#   Permission is hereby granted, free of charge, to any person obtaining a copy
#   of this software and associated documentation files (the "Software"), to deal
#   in the Software without restriction, including without limitation the rights
#   to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
#   copies of the Software, and to permit persons to whom the Software is
#   furnished to do so, subject to the following conditions:
#   
#   The above copyright notice and this permission notice shall be included in
#   all copies or substantial portions of the Software.
#   
#   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#   OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
#   THE SOFTWARE.
#################################################################################

import sys, os
from collections import defaultdict


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Functions        


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
from collections import namedtuple

t_locus = namedtuple("t_locus", "contig strand beg end")
def size (self):
    return self.end - self.beg

def from_str (d):
    return t_locus(str(d[0]), str(d[1]), int(d[2]), int(d[3]))


def overlaps (*loci):
    """
    takes loci in [chrm, beg, end] format and 
    returns if overlap
    """
    # same contigs?
    if  (loci[0].contig != loci[1].contig or
         loci[0].strand != loci[1].strand):
        return False

    # does first loci end extend past second loci beg?
    loci = sorted(loci)
    return loci[0].end > loci[1].beg
    
    
t_locus.size = size
t_locus.overlaps = overlaps
t_locus.from_str = staticmethod(from_str)

        
#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   t_loci


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
#----------------------------------------------------------------
class t_loci (object):
    def __init__ (self, other_loci=None):
        if other_loci == None:
            self.data = []
        if isinstance(other_loci, t_loci):
            self.data = list(other.loci)
        else:
            self.data = []
            for l in other_loci:
                if not isinstance(l, t_locus):
                    raise Exception("t_loci.__init__() expects a sequence of t_locus")
                self.data.append(l)

    def __repr__ (self):
        return json.dumps(self.data, cls=t_dump_locus)

    def __cmp__ (self, other):
        return cmp(self.data, other.data)


    def merge_neighbours(self, min_distance ):
        """
            combine_close_neighbours

                Merge neighbours that are less than a certain
                distance apart.

                side effect: sort 

        """
        if not len(self.data): return

        new_data = []

        self.data.sort()

        curr_contig, curr_strand, curr_beg, curr_end = self.data[0]

        for l in self.data[1:]:

            if  (curr_contig != l.contig or 
                 curr_strand != l.strand or
                 l.beg - curr_end + 1 > min_distance):
                new_data.append( t_locus(curr_contig, curr_strand, curr_beg, curr_end) )
                curr_contig, curr_strand, curr_beg, curr_end = l
                continue

            curr_end = max(l.end, curr_end)

        new_data.append( t_locus(curr_contig, curr_strand, curr_beg, curr_end) )

        self.data = new_data

        return self


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   json

#   json.dumps(loci, cls=t_dump_locus)                   
#   json.loads(saved_values, object_hook = t_load_locus) 


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
def t_load_locus(dct):
    if '__t_locus__' in dct:
        return t_locus(dct['contig'], dct['strand'], dct['beg'], dct['end'])
    elif '__t_loci__' in dct:
        return t_locus(dct['contig'], dct['strand'], dct['beg'], dct['end'])
    return dct

import json
class t_dump_locus(json.JSONEncoder):
    def _iterencode(self, o, markers=None):
        if isinstance(o, t_locus):
            d = o._asdict()
            d['__t_locus__'] =True
            return json.JSONEncoder._iterencode(self, d, markers)
        return json.JSONEncoder._iterencode(self, o, markers)





#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Testing


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
import unittest
import json
def as_t_locus(dct):
    if '__t_locus__' in dct:
        return t_locus( str(dct['contig']), 
                        int(dct['beg']), 
                        int(dct['end']), 
                        str(dct['strand']))
    return dct

class Test_locus(unittest.TestCase):

    #       self.assertEqual(self.seq, range(10))
    #       self.assert_(element in self.seq)
    #       self.assertRaises(ValueError, random.sample, self.seq, 20)



    def test_t_locus(self):
        """
            test 
        """
        test_locus = t_locus("chr1", "+", 1, 3) 
        loci = [
                t_locus("chr2", "-", 44, 53, ),
                t_locus("chr1", "+", 2, 5,   ),
                t_locus("chr1", "+", 3, 5,   ),
                ]
            
        self.assertEqual(any(l.overlaps(test_locus) for l in loci), True)
        self.assertEqual(all(l.overlaps(test_locus) for l in loci), False)
        
        saved_values = json.dumps(loci)
        new_loci = json.loads(saved_values)
        new_loci = [t_locus.from_str(l) for l in new_loci]
        self.assertEqual(new_loci, loci)
        
        
        
                
        saved_values = json.dumps(loci, cls=t_dump_locus)
        new_loci     = json.loads(saved_values, object_hook = t_load_locus)
        self.assertEqual(new_loci, loci)
        
    def test_t_loci_merge_neighbours(self):
        """
            test 
        """
        test_locus = t_locus("chr1", "+", 1, 3) 
        loci = [
                t_locus("chr2", "-", 44, 53, ),
                t_locus("chr1", "+", 2, 5,   ),
                t_locus("chr1", "-", 3, 5,   ),
                t_locus("chr1", "+", 10, 11,   ),
                ]

        
        orig_loci = t_loci(loci)        
        print json.dumps(orig_loci, indent = 4, cls=t_dump_locus)
        orig_loci.merge_neighbours(20)
        print json.dumps(orig_loci, indent = 4, cls=t_dump_locus)

#
#   debug code not run if called as a module
#     
if __name__ == '__main__':
    if sys.argv.count("--debug"):
        sys.argv.remove("--debug")
    unittest.main()



    
